package com.lxq.cloud.ai.service;

import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

import java.util.Map;

@Service
@Slf4j
public class ChatServiceImpl implements ChatService{


    private final ChatModel chatModel;

    public ChatServiceImpl(ChatModel chatModel) {
        this.chatModel = chatModel;
    }


    @Override
    public String normalCompletion(String input) {
        Prompt prompt = new Prompt(input);
        ChatResponse response = chatModel.call(prompt);
        String content = response.getResult().getOutput().getText();
        log.info("response: {}", response);
        return content;
    }


    @Override
    public Map<String, String> streamCompletion(String message) {


        StringBuilder fullContent = new StringBuilder();


        chatModel.stream(new Prompt(message))

                .flatMap(chatResponse -> Flux.fromIterable(chatResponse.getResults()))

                .map(content -> content.getOutput().getText())

                .doOnNext(fullContent::append)

                .last()

                .map(lastContent -> Map.of(message, fullContent.toString()))

                .block();
        String response = fullContent.toString();
        log.info("response: {}", response);
        return Map.of(message, response);

    }


}